from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import List, Optional
import uvicorn
import os

# 强制使用离线模式
os.environ['LLM_MODE'] = 'local'
os.environ['LOCAL_MODEL_NAME'] = 'offline'

from services.llm_service import LLMService
from services.rag_service import RAGService
from services.arxiv_service import ArxivService

# 创建FastAPI应用
app = FastAPI(title="论文搜索引擎API", version="1.0.0")

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 初始化服务
llm_service = LLMService()
rag_service = RAGService()
arxiv_service = ArxivService()

# 数据模型
class SearchRequest(BaseModel):
    query: str
    limit: int = 10
    filters: Optional[dict] = None

class ChatRequest(BaseModel):
    question: str
    context_papers: Optional[List[str]] = None
    chat_history: Optional[List[dict]] = None

class Paper(BaseModel):
    id: str
    title: str
    authors: List[str]
    abstract: str
    arxiv_id: Optional[str] = None
    categories: List[str]
    relevance_score: float

class ChatResponse(BaseModel):
    answer: str
    sources: List[dict]
    confidence: float

@app.on_event("startup")
async def startup_event():
    """应用启动时初始化服务"""
    try:
        # 初始化LLM服务
        await llm_service.initialize()
        
        # 初始化RAG服务
        await rag_service.initialize()
        
        print("✅ 所有服务初始化完成")
    except Exception as e:
        print(f"❌ 服务初始化失败: {e}")

@app.get("/")
async def root():
    """根端点"""
    return {
        "message": "论文搜索引擎API服务运行中",
        "version": "1.0.0",
        "status": "active"
    }

@app.get("/health")
async def health_check():
    """健康检查端点"""
    return {
        "status": "healthy",
        "llm_ready": llm_service.is_ready(),
        "rag_ready": rag_service.is_ready()
    }

@app.post("/search", response_model=List[Paper])
async def search_papers(request: SearchRequest):
    """搜索论文"""
    try:
        # 使用Arxiv服务搜索论文
        papers = await arxiv_service.search_papers(
            query=request.query,
            max_results=request.limit,
            filters=request.filters
        )
        return papers
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"搜索失败: {str(e)}")

@app.post("/chat", response_model=ChatResponse)
async def chat_with_ai(request: ChatRequest):
    """与AI聊天（基于RAG）"""
    try:
        if not rag_service.is_ready():
            raise HTTPException(status_code=503, detail="RAG服务未就绪")
        
        response = await rag_service.generate_answer(
            question=request.question,
            context_papers=request.context_papers,
            chat_history=request.chat_history
        )
        return response
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"聊天失败: {str(e)}")

@app.get("/paper/{paper_id}/download")
async def download_paper(paper_id: str):
    """下载论文"""
    try:
        # 这里可以实现论文下载逻辑
        # 目前返回论文信息
        paper_info = await arxiv_service.get_paper_info(paper_id)
        return {
            "paper_id": paper_id,
            "download_url": f"https://arxiv.org/pdf/{paper_id}",
            "info": paper_info
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"下载失败: {str(e)}")

@app.get("/statistics")
async def get_statistics():
    """获取统计信息"""
    try:
        stats = await rag_service.get_statistics()
        return stats
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取统计信息失败: {str(e)}")

if __name__ == "__main__":
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=8000,
        reload=True,
        log_level="info"
    )